# Please install OpenAI SDK first: `pip3 install openai`

import json

from openai import OpenAI


# 定义一个函数，用于根据日期、出发地和目的地查询航班号
def get_flight_number(date: str, departure: str, destination: str):
    """
    查询航班号
    :param date: 日期
    :param departure: 出发地
    :param destination: 目的地
    :return: 航班号
    """
    # 定义一个嵌套字典，存储不同出发地和目的地对应的航班号
    flight_number = {
        "北京": {
            "上海": "1234",
            "广州": "8321",
        },
        "上海": {
            "北京": "1233",
            "广州": "8123",
        }
    }
    # 返回查询到的航班号-方法的核心逻辑
    destination_data = flight_number[departure][destination]
    return {"flight_number": destination_data}


# 定义一个函数，用于根据日期和航班号查询票价
def get_ticket_price(date: str, flight_number: str):
    """
    查询票价
    :param date: 日期
    :param flight_number: 航班号
    :return: 票价
    """
    # 目前代码中票价是固定的，返回一个固定的票价值
    return {"ticket_price": "1000"}


# 定义一个列表，存储工具函数的定义
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_flight_number",
            "description": "根据始发地、目的地和日期，查询对应日期的航班号",
            "parameters": {
                "type": "object",
                "properties": {
                    "departure": {
                        "description": "出发地",
                        "type": "string"
                    },
                    "destination": {
                        "description": "目的地",
                        "type": "string"
                    },
                    "date": {
                        "description": "日期",
                        "type": "string",
                    }
                },
                "required": ["departure", "destination", "date"]
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_ticket_price",
            "description": "查询某航班在某日的票价",
            "parameters": {
                "type": "object",
                "properties": {
                    "flight_number": {
                        "description": "航班号",
                        "type": "string"
                    },
                    "date": {
                        "description": "日期",
                        "type": "string",
                    }
                },
                "required": ["flight_number", "date"]
            },
        }
    },
]

# 初始化DeepSeek客户端，填写自己的APIKey
client = OpenAI(api_key="sk-33fce941cea3409484cdbe7a9abca430", base_url="https://api.deepseek.com")


# 定义一个函数，用于解析模型的函数调用结果
def parse_function_call(model_response, messages):
    """
    解析模型的函数调用结果
    :param model_response: 模型返回的响应
    :param messages: 消息列表
    """
    # 判断模型是否进行了函数调用
    if model_response.choices[0].message.tool_calls:
        # 获取函数调用信息
        tool_call = model_response.choices[0].message.tool_calls[0]
        # 获取函数调用的参数
        args = tool_call.function.arguments
        # 初始化函数调用结果
        function_result = {}
        # 根据函数名称调用对应的函数
        if tool_call.function.name == "get_flight_number":
            # 调用get_flight_number函数，并将结果存储到function_result中
            function_result = get_flight_number(**json.loads(args))
        if tool_call.function.name == "get_ticket_price":
            # 调用get_ticket_price函数，并将结果存储到function_result中
            function_result = get_ticket_price(**json.loads(args))
        # 构造tool message，将函数调用结果添加到消息列表中
        messages.append({
            "role": "tool",
            "content": f"{json.dumps(function_result)}",
            "tool_call_id": tool_call.id
        })
        # 再次调用模型，将函数调用结果输入模型
        response = client.chat.completions.create(
            model="deepseek-chat",  # 填写需要调用的模型名称
            messages=messages,
            tools=tools,
        )
        # 打印最终模型的回答结果
        print("------>第二次调用模型的请求参数：", json.dumps(messages, ensure_ascii=False))
        print("------>第二次调用模型的回答结果：", response.choices[0].message.model_dump_json())
        # 将最终模型的回答结果添加到消息列表中
        messages.append(response.choices[0].message.model_dump())


# 清空对话消息列表
messages = []
# 拼接第一次对话，添加系统消息和用户消息
messages.append(
    {"role": "system", "content": "不要假设或猜测传入函数的参数值。如果用户的描述不明确，请要求用户提供必要信息"})
messages.append({"role": "user", "content": "帮我查询1月23日，北京到广州的航班"})
# 调用模型处理用户消息
response = client.chat.completions.create(
    model="deepseek-chat",  # 填写需要调用的模型名称
    messages=messages,
    tools=tools,
)
# 打印模型的响应
print("------>第一次调用模型的请求参数：", json.dumps(messages, ensure_ascii=False))
print("------>第一次调用模型的回答结果：", response.choices[0].message.model_dump_json())
# 将模型的响应添加到消息列表中
messages.append(response.choices[0].message.model_dump())
# 调用parse_function_call函数，解析函数调用结果
parse_function_call(response, messages)
